package main

import (
	"fmt"
)

func main() {
	//a := 0
	//fmt.Scan(&a)
	//fmt.Printf("%d\n", a)

	// 测试用例
	records := [][]int{
		{1, 2},
		{2, 3},
		{3, 5},
		{4, 5},
		{2, 6},
		{6, 8},
	}

	treeMap := GenerateTreeNodeBySourceTarget(records)

	searchNodeVal := 2
	upstreams, downStreams := findUpstreamTreeNode(treeMap, searchNodeVal)

	fmt.Printf("upstreams as below:\n")
	for _, num := range upstreams {
		fmt.Printf("%v, ", num)
	}

	fmt.Printf("\n")

	fmt.Printf("downstreams as below:\n")
	for _, num := range downStreams {
		fmt.Printf("%v, ", num)
	}
}

type TreeNode struct {
	ParentValSet  map[int]struct{} // 父子数组去重
	SubNodeValSet map[int]struct{}
	ParentNodes   []*TreeNode
	SubNodes      []*TreeNode
	Val           int
}

func GenerateTreeNodeBySourceTarget(sourceTargets [][]int) (searchMap map[int]*TreeNode) { // 图 数据结构，暂时不考虑根节点
	// 默认sourceTargets 每个Item中都是两个Idx，表示Source 和 Target
	searchMap = make(map[int]*TreeNode)

	for i := 0; i < len(sourceTargets); i++ {
		sourceVal := sourceTargets[i][0]
		targetVal := sourceTargets[i][1]

		// 开始考虑构建节点及Map, 如果节点不存在，则先初始化
		var sourceNode, targetNode *TreeNode
		if _, has := searchMap[sourceVal]; !has {
			searchMap[sourceVal] = &TreeNode{
				ParentValSet:  make(map[int]struct{}),
				SubNodeValSet: make(map[int]struct{}),
				Val:           sourceVal,
			}
		}
		if _, has := searchMap[targetVal]; !has {
			searchMap[targetVal] = &TreeNode{
				ParentValSet:  make(map[int]struct{}),
				SubNodeValSet: make(map[int]struct{}),
				Val:           targetVal,
			}
		}
		sourceNode = searchMap[sourceVal]
		targetNode = searchMap[targetVal]

		// 构建父子关系
		if _, has := sourceNode.SubNodeValSet[targetVal]; !has {
			sourceNode.SubNodes = append(sourceNode.SubNodes, targetNode)
			sourceNode.SubNodeValSet[targetVal] = struct{}{}
		}

		if _, has := targetNode.ParentValSet[sourceVal]; !has {
			targetNode.ParentNodes = append(targetNode.ParentNodes, sourceNode)
			targetNode.ParentValSet[sourceVal] = struct{}{}
		}
	}

	return searchMap
}

func findUpstreamTreeNode(searchMap map[int]*TreeNode, nodeVal int) (upStreams, downStreams []int) {
	// 要求：快速生成一棵树，找出随意一个节点的上游
	// 思路1: 如果要生成树，那树节点需要记录自己的parent Node数组和sub Node数组
	//       要快速找到任意一个节点，最简单的思路就是使用Map存储节点和Tree Node的映射

	if _, has := searchMap[nodeVal]; !has {
		return
	}

	// 构建结果
	currNode := searchMap[nodeVal]

	for _, node := range currNode.ParentNodes {
		upStreams = append(upStreams, node.Val)
	}

	for _, node := range currNode.SubNodes {
		downStreams = append(downStreams, node.Val)
	}

	return
}
